{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 385
    },
    "id": "smwr7jBZfEqX",
    "outputId": "034aa1c3-86a8-4ad9-ce52-e2f0d836ee56"
   },
   "outputs": [],
   "source": [
    "from __future__ import division\n",
    "import numpy as np\n",
    "import cupy as cp\n",
    "from keras.datasets import mnist \n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "G3yLT___dPYB"
   },
   "outputs": [],
   "source": [
    "##### parameter setting #####\n",
    "cp.random.seed(0)\n",
    "thr = [100, 100] \n",
    "lrF = [0.01, 0.01]  \n",
    "lr = [0.1, 0.1]\n",
    "b = [5, 48]  \n",
    "a = [0, 0]\n",
    "lamda = [0, 0] \n",
    "dropout = [0, 0]  \n",
    "NumOfCatgry = 10\n",
    "Tend = 255\n",
    "Tmax = 256\n",
    "Size_i = 28\n",
    "N_ch = 1\n",
    "fact = [5, 5]  \n",
    "NhidenNeurons = 600\n",
    "Num_neurons  = 784\n",
    "Nepoch = 200\n",
    "gamma = 1\n",
    "best = 0  \n",
    "Trbest = 0\n",
    "th=[]\n",
    "to=[]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Mva5rnvrZ55e"
   },
   "outputs": [],
   "source": [
    "#### data prepration ####\n",
    "\n",
    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "\n",
    "Images = x_train\n",
    "Labels = y_train\n",
    "Images = np.array(Images)\n",
    "images = []\n",
    "labels = []\n",
    "for i in range(len(Labels)):\n",
    "    if Labels[i] < NumOfCatgry :\n",
    "        images.append(np.floor(Tend-(np.dot(Images[i],Tend)/255).reshape(1,N_ch,Size_i,Size_i))) \n",
    "        labels.append(Labels[i])          \n",
    "\n",
    "\n",
    "Images = x_test\n",
    "Labels = y_test\n",
    "Images = np.array(Images)\n",
    "images_test = []\n",
    "labels_test = []\n",
    "for i in range(len(Labels)):\n",
    "    if Labels[i] < NumOfCatgry:\n",
    "        images_test.append(np.floor(Tend-(np.dot(Images[i],Tend)/255).reshape(1,N_ch,Size_i,Size_i)))\n",
    "        labels_test.append(Labels[i])\n",
    "\n",
    "\n",
    "\n",
    "images = cp.asarray(images)\n",
    "images_test = cp.asarray(images_test)\n",
    "labels = cp.asarray(labels)\n",
    "labels_test = cp.asarray(labels_test)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vA-TlZ2AdbyM"
   },
   "outputs": [],
   "source": [
    "\n",
    "############## Define forward function for Fully Connected layers: ##############\n",
    "def Forward_FC(FT, weight, Num_in, Num_out, init, t):\n",
    "    if (init==1):\n",
    "        t = cp.mgrid[0:Num_in, 0:1]\n",
    "        Spikes = cp.zeros((Num_in, 1, Tmax + 1))\n",
    "    else:\n",
    "        t = t\n",
    "        Spikes = cp.zeros((Num_in, 1, Tmax + 1))\n",
    "\n",
    "    thre = thr[layer]/fact[layer]\n",
    "    Spikes[t[0], t[1], FT.astype(int)] = 1    \n",
    "    Voltage = cp.cumsum(cp.tensordot(cp.sign(weight), Spikes), 1)  \n",
    "    Voltage[:,Tmax] = thre + 1\n",
    "    ST = (cp.argmax(Voltage>thre,axis=1).astype(float)+1).reshape(Num_out,1) \n",
    "    ST[ST>Tmax] = Tmax\n",
    "\n",
    "    return (ST, Voltage, t)\n",
    "\n",
    "############## Define target function: ##############\n",
    "def target_function(inp, Vinp, correct):\n",
    "\n",
    "    target = cp.zeros([NumOfCatgry,1])\n",
    "    winner = cp.argmin(inp)\n",
    "    Vwin = cp.argmax(Vinp[:,inp[winner].astype(int)])        \n",
    "    EQ = inp==inp[winner]  \n",
    "    temp = inp[winner]\n",
    "\n",
    "    minFiring = min(inp)             \n",
    "\n",
    "    if minFiring == Tmax:\n",
    "        target[:] = minFiring\n",
    "        target[labels[IMGs[iteration]]] = minFiring \n",
    "        target = target.astype(int)\n",
    "    else:\n",
    "        target[:] = inp\n",
    "        toChange = (inp - minFiring) < gamma\n",
    "        target[toChange] = min(minFiring + gamma, Tmax)\n",
    "        target[labels[IMGs[iteration]]] = minFiring\n",
    "\n",
    "    if minFiring==Tmax:\n",
    "        V=cp.argmax(Vinp[:,Tmax-3])\n",
    "        if V==labels[IMGs[iteration]]:\n",
    "          correct+=1\n",
    "    else:\n",
    "        if inp[labels[IMGs[iteration]]]==minFiring:\n",
    "            correct+=1\n",
    "\n",
    "    return(target,correct)\n",
    "\n",
    "\n",
    "############## Dropout function: ##############\n",
    "def DP(dp, inp, Num_in):\n",
    "    if (dp>0):\n",
    "        inp[cp.asarray(cp.random.permutation(Num_in)[:dp])] = Tmax+1  \n",
    "    return(inp)\n",
    "\n",
    "############## Training of the Fully Connected layers: ##############\n",
    "def backward_FC(delta, inp, out, weight, lr, lamda):        \n",
    "    HF = (inp<out) \n",
    "    weight -= (lr*delta)[:,cp.newaxis]*HF \n",
    "    weight -= lr*lamda*weight\n",
    "    fact[layer] -= (lrF[layer]*delta[:, cp.newaxis]*HF*cp.sign(weight)).mean() \n",
    "    Delta = cp.sum((weight*delta[:,cp.newaxis])*HF,0)    \n",
    "    return(weight,Delta)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rcw_V33ddfFr"
   },
   "outputs": [],
   "source": [
    "##### initialization #####\n",
    "print('Do yo want to load and test the pretrained model? (y/n)')\n",
    "pretrain= input()\n",
    "if pretrain=='y':\n",
    "    W0 = cp.asarray(np.load('./W0_binary.npy'))\n",
    "    W1 = cp.asarray(np.load('./W1_binary.npy'))\n",
    "    fact = [4.03707844, 9.06516109]\n",
    "    eposchs=0\n",
    "else:\n",
    "    layer = 0\n",
    "    W0 = (b[layer] - a[layer]) * cp.random.random_sample((NhidenNeurons,Num_neurons,1)) + a[layer]\n",
    "    layer = 1\n",
    "    W1 = (b[layer] - a[layer]) * cp.random.random_sample((NumOfCatgry,NhidenNeurons,1)) + a[layer]\n",
    "    eposchs=Nepoch\n",
    "\n",
    "\n",
    "#### Training ####\n",
    "init=1\n",
    "for epoch in range(eposchs):\n",
    "    correct=0.0\n",
    "    IMGs = cp.random.randint(0, 60000, 60000, dtype=int)  \n",
    "    for iteration in tqdm(range(len(IMGs)), position=0):   \n",
    "\n",
    "        #### Forward processing: #####\n",
    "        layer = 0               \n",
    "        FiringTime0 = cp.reshape(images[IMGs[iteration]],[Num_neurons,1])        \n",
    "        (FiringTime1,VoltageH,th) = Forward_FC(FiringTime0, W0, Num_neurons, NhidenNeurons, init, th)\n",
    "        layer = 1  \n",
    "        (SpikeTime,VoltageO,to) = Forward_FC(FiringTime1, W1, NhidenNeurons, NumOfCatgry, init, to)\n",
    "   \n",
    "        ######## error and target time: ####\n",
    "        (target, correct) = target_function(SpikeTime, VoltageO, correct)\n",
    "\n",
    "        ##### Backward processing (Training): ######\n",
    "        layer= 1   \n",
    "        delta3 = cp.zeros([NumOfCatgry,1])       \n",
    "        delta3 = ((target-SpikeTime)/Tmax)  \n",
    "        norm = cp.linalg.norm(delta3)\n",
    "        if (norm!=0):\n",
    "            delta3 = delta3/norm       \n",
    "        FiringTime1 = DP(dropout[layer-1], FiringTime1, NhidenNeurons)\n",
    "        (W1,delta2) = backward_FC(delta3,FiringTime1,SpikeTime[:,cp.newaxis],W1,lr[layer],lamda[layer]) \n",
    "\n",
    "        layer= 0     \n",
    "        norm = cp.linalg.norm(delta2)\n",
    "        if (norm!=0):\n",
    "            delta2 = delta2/norm       \n",
    "        FiringTime0 = DP(dropout[layer-1], FiringTime0, Num_neurons) \n",
    "        (W0,delta1) = backward_FC(delta2,FiringTime0,FiringTime1[:,cp.newaxis],W0,lr[layer],lamda[layer]) \n",
    "       \n",
    "        init = 0                 \n",
    "\n",
    " \n",
    "        ######## TEST: ########\n",
    "    correct=0\n",
    "    for iteration in tqdm(range(len(images_test)), position=0):\n",
    "        ############# Forward pass: ##############\n",
    "        layer = 0               \n",
    "        FiringTime0 = cp.reshape(images_test[iteration],[Num_neurons,1])        \n",
    "        (FiringTime1,VoltageH,th) = Forward_FC(FiringTime0, W0, Num_neurons, NhidenNeurons, init, th)\n",
    "\n",
    "        layer = 1  \n",
    "        (SpikeTime,VoltageO,to) = Forward_FC(FiringTime1, W1, NhidenNeurons, NumOfCatgry, init, to)\n",
    "\n",
    "        ############## Decision making ##############\n",
    "        winner = cp.argmin(SpikeTime)\n",
    "        Vwin = cp.argmax(VoltageO[:,SpikeTime[winner].astype(int)])        \n",
    "        EQ = SpikeTime==SpikeTime[winner]  \n",
    "        temp = SpikeTime[winner]\n",
    "        minFiring = min(SpikeTime)\n",
    "        maxFiring = max(SpikeTime)             \n",
    "\n",
    "        if minFiring==Tmax:\n",
    "            V=np.argmax(VoltageO[:,Tmax-3])\n",
    "            if V==labels_test[iteration]:\n",
    "              correct+=1\n",
    "        else:\n",
    "            if SpikeTime[labels_test[iteration]]==minFiring:\n",
    "                correct+=1             \n",
    "\n",
    "    print('epoch = ', epoch, 'Test accuracy= ', correct/len(images_test), 'Best Accuracy= ', best)\n",
    "\n",
    "    if (correct/len(images_test)>=best):\n",
    "        best = correct/len(images_test)\n",
    "        np.save('./W0_binary_best.npy',cp.asnumpy(W0))\n",
    "        np.save('./W1_binary_best.npy',cp.asnumpy(W1))\n",
    "    \n",
    "    \n",
    "######## Fibnal TEST result: ########\n",
    "correct=0\n",
    "for iteration in tqdm(range(len(images_test)), position=0):#progressbar.progressbar(range(len(images_test))):   #for iteration in range(len(images_test)): \n",
    "    ############# Forward pass: ##############\n",
    "    layer = 0               \n",
    "    FiringTime0 = cp.reshape(images_test[iteration],[Num_neurons,1])        \n",
    "    (FiringTime1,VoltageH,th) = Forward_FC(FiringTime0, W0, Num_neurons, NhidenNeurons, init, th)\n",
    "\n",
    "    layer = 1  \n",
    "    (SpikeTime,VoltageO,to) = Forward_FC(FiringTime1, W1, NhidenNeurons, NumOfCatgry, init, to)\n",
    "\n",
    "    ############## Decision making ##############\n",
    "    winner = cp.argmin(SpikeTime)\n",
    "    Vwin = cp.argmax(VoltageO[:,SpikeTime[winner].astype(int)])        \n",
    "    EQ = SpikeTime==SpikeTime[winner]  \n",
    "    temp = SpikeTime[winner]\n",
    "    minFiring = min(SpikeTime)\n",
    "    maxFiring = max(SpikeTime)             \n",
    "\n",
    "    if minFiring==Tmax:\n",
    "        V=np.argmax(VoltageO[:,Tmax-3])\n",
    "        if V==labels_test[iteration]:\n",
    "          correct+=1\n",
    "    else:\n",
    "        if SpikeTime[labels_test[iteration]]==minFiring:\n",
    "            correct+=1             \n",
    "\n",
    "print('Final Test accuracy= ', correct/len(images_test))\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "BS4NN_mnist.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "ganomaly",
   "language": "python",
   "name": "ganomaly"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
